import torch
import torch.nn as nn
import torch.nn.functional as F


class DenseRelationDistill(nn.Module):

    def __init__(self, indim, keydim, valdim, dense_sum=False):
        super(DenseRelationDistill, self).__init__()
        # self.key_q = nn.Conv2d(indim, keydim, kernel_size=(3,3), padding=(1,1), stride=1)
        # self.value_q = nn.Conv2d(indim, valdim, kernel_size=(3,3), padding=(1,1), stride=1)
        self.key_t = nn.Conv2d(indim, keydim, kernel_size=(3, 3), padding=(1, 1), stride=1)
        self.value_t = nn.Conv2d(indim, valdim, kernel_size=(3, 3), padding=(1, 1), stride=1)
        self.sum = dense_sum
        if self.sum:
            self.key_q0 = nn.Conv2d(indim, keydim, kernel_size=(3, 3), padding=(1, 1), stride=1)
            self.value_q0 = nn.Conv2d(indim, valdim, kernel_size=(3, 3), padding=(1, 1), stride=1)
            self.key_q1 = nn.Conv2d(indim, keydim, kernel_size=(3, 3), padding=(1, 1), stride=1)
            self.value_q1 = nn.Conv2d(indim, valdim, kernel_size=(3, 3), padding=(1, 1), stride=1)
            self.key_q2 = nn.Conv2d(indim, keydim, kernel_size=(3, 3), padding=(1, 1), stride=1)
            self.value_q2 = nn.Conv2d(indim, valdim, kernel_size=(3, 3), padding=(1, 1), stride=1)
            self.key_q3 = nn.Conv2d(indim, keydim, kernel_size=(3, 3), padding=(1, 1), stride=1)
            self.value_q3 = nn.Conv2d(indim, valdim, kernel_size=(3, 3), padding=(1, 1), stride=1)
            self.key_q4 = nn.Conv2d(indim, keydim, kernel_size=(3, 3), padding=(1, 1), stride=1)
            self.value_q4 = nn.Conv2d(indim, valdim, kernel_size=(3, 3), padding=(1, 1), stride=1)
            self.bnn0 = nn.BatchNorm2d(256)
            self.bnn1 = nn.BatchNorm2d(256)
            self.bnn2 = nn.BatchNorm2d(256)
            self.bnn3 = nn.BatchNorm2d(256)
            self.bnn4 = nn.BatchNorm2d(256)
            self.combine = nn.Conv2d(512, 256, kernel_size=1, padding=0, stride=1)

    def forward(self, features, attentions):
        features = list(features)
        if isinstance(attentions, dict):
            for i in range(len(attentions)):
                if i == 0:
                    atten = attentions[i].unsqueeze(0)
                else:
                    atten = torch.cat((atten, attentions[i].unsqueeze(0)), dim=0)
            attentions = atten.cuda()
        output = []
        h, w = attentions.shape[2:]
        ncls = attentions.shape[0]
        key_t = self.key_t(attentions)
        val_t = self.value_t(attentions)
        for idx in range(len(features)):
            feature = features[idx]
            bs = feature.shape[0]
            H, W = feature.shape[2:]
            feature = F.interpolate(feature, size=(h, w), mode='bilinear', align_corners=True)
            key_q = eval('self.key_q' + str(idx))(feature).view(bs, 32, -1)
            val_q = eval('self.value_q' + str(idx))(feature)
            for i in range(bs):
                kq = key_q[i].unsqueeze(0).permute(0, 2, 1)
                vq = val_q[i].unsqueeze(0)

                p = torch.matmul(kq, key_t.view(ncls, 32, -1))
                p = F.softmax(p, dim=1)

                val_t_out = torch.matmul(val_t.view(ncls, 128, -1), p).view(ncls, 128, h, w)
                for j in range(ncls):
                    if (j == 0):
                        final_2 = torch.cat((vq, val_t_out[j].unsqueeze(0)), dim=1)
                    else:
                        final_2 += torch.cat((vq, val_t_out[j].unsqueeze(0)), dim=1)
                if (i == 0):
                    final_1 = final_2
                else:
                    final_1 = torch.cat((final_1, final_2), dim=0)
            final_1 = F.interpolate(final_1, size=(H, W), mode='bilinear', align_corners=True)
            if self.sum:
                final_1 = eval('self.bnn' + str(idx))(final_1)

            output.append(final_1)

        if self.sum:
            for i in range(len(output)):
                output[i] = self.combine(torch.cat((features[i], output[i]), dim=1))
        output = tuple(output)

        return output


# 注释
class NoUseDenseRelationDistill(DenseRelationDistill):
    def __init__(self, indim, keydim, valdim, dense_sum=False):
        super(NoUseDenseRelationDistill, self).__init__(indim, keydim, valdim, dense_sum)
        # DenseRelationDistill(256,32,128,self.dense_sum)

    def forward(self, features, attentions):
        features = list(features)
        # len features 5 代表5层来自不同stage的不同分辨率的特征图
        if isinstance(attentions, dict):
            for i in range(len(attentions)):
                if i == 0:
                    atten = attentions[i].unsqueeze(0)
                else:
                    atten = torch.cat((atten, attentions[i].unsqueeze(0)), dim=0)
            attentions = atten.cuda()
        output = []
        # batch size 4
        # attentions shape: 20,256,16,16
        # 注意力特征的大小固定，都是16x16
        h, w = attentions.shape[2:]
        ncls = attentions.shape[0]
        # ncls: 20
        key_t = self.key_t(attentions)
        val_t = self.value_t(attentions)
        # key_t shape 20,32,16,16
        # val_t shape 20 128 16 16
        for idx in range(len(features)):
            feature = features[idx]
            # 查询特征的大小不固定
            # feature shape 4,256,272,304
            # feature shape 4,256,136,152
            bs = feature.shape[0]
            H, W = feature.shape[2:]
            # H,W=272,304
            # h,w=16,16
            feature = F.interpolate(feature, size=(h, w), mode='bilinear', align_corners=True)
            # 5种尺寸的特征图对应于5个查询编码器
            key_q = eval('self.key_q' + str(idx))(feature).view(bs, 32, -1)
            # key_q shape 4,32,256
            # key_q: bsx32xhw
            val_q = eval('self.value_q' + str(idx))(feature)
            # val_q shape 4,128,16,16
            for i in range(bs):
                kq = key_q[i].unsqueeze(0).permute(0, 2, 1)
                # kq shape 1,256,32
                # kq: 1 x hw x 32
                vq = val_q[i].unsqueeze(0)
                # vq shape 1,128,16,16
                # key_t.view: Nx32(C/8)xhw
                # key_t.view: 20x32(C/8)x256
                p = torch.matmul(kq, key_t.view(ncls, 32, -1))
                p = F.softmax(p, dim=1)
                # p shape 20,256,256
                # p: N x hw x hw
                # val_t 注意力
                # 支持值跟权重相乘之后与查询值进行拼接，支持值加权
                # val_t.view 20,128,256
                val_t_out = torch.matmul(val_t.view(ncls, 128, -1), p).view(ncls, 128, h, w)
                # val_t_out shape 20,128,16,16
                # val_t_out: N x 128 x h x w
                for j in range(ncls):
                    if j == 0:
                        final_2 = torch.cat((vq, val_t_out[j].unsqueeze(0)), dim=1)
                    else:
                        final_2 += torch.cat((vq, val_t_out[j].unsqueeze(0)), dim=1)
                # 20个支持值映射加权求和
                # final_2: 1, 256 x h x w
                # final_2: 1, 256, 16, 16
                if i == 0:
                    final_1 = final_2
                else:
                    final_1 = torch.cat((final_1, final_2), dim=0)

            final_1 = F.interpolate(final_1, size=(H, W), mode='bilinear', align_corners=True)
            # final_1: bs x 256 x h x w
            # final_1: 4,256,272,304
            if self.sum:
                final_1 = eval('self.bnn' + str(idx))(final_1)

            output.append(final_1)

        if self.sum:
            for i in range(len(output)):
                output[i] = self.combine(torch.cat((features[i], output[i]), dim=1))
        output = tuple(output)

        return output
